package com.web.logindemo.aspectj;


import com.google.common.collect.Maps;
import com.web.logindemo.common.annotation.RateLimiter;
import com.web.logindemo.common.enums.LimitType;
import com.web.logindemo.common.exception.ServiceException;
import com.web.logindemo.common.util.ServletUtils;
import com.web.logindemo.common.util.StringUtils;
import com.web.logindemo.common.util.ip.IPUtils;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.net.SocketException;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * 限流处理
 */
@Aspect
@Component
public class RateLimiterAspect {

    private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);

    private RedisTemplate<Object, Object> redisTemplate;

    private RedisScript limitScript;

    /**
     * 本地限流兜底 将 key存在map中
     *
     * @param redisTemplate
     */
    private final Map<String, com.google.common.util.concurrent.RateLimiter> limiterMap = Maps.newConcurrentMap();


    @Autowired
    public void setRedisTemplate1(RedisTemplate<Object, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Autowired
    public void setLimitScript(RedisScript<Long> limitScript) {
        this.limitScript = limitScript;
    }

    // 配置织入点
    @Pointcut("@annotation(com.web.logindemo.common.annotation.RateLimiter)")
    public void rateLimiterPointCut() {

    }

    @Before("rateLimiterPointCut()")
    public void doBefore(JoinPoint joinPoint) {
        RateLimiter rateLimiter = getAnnotationRateLimiter(joinPoint);
        String key = rateLimiter.key();
        int time = rateLimiter.time();
        int count = rateLimiter.count();
        String combineKey = getCombineKey(rateLimiter, joinPoint);
        List<Object> keys = Collections.singletonList(combineKey);
        try {

//            boolean isRedisConnection = redisTemplate.sort();
//            if (!isRedisConnection) {
//                // redis宕机 增加兜底策略
//                rateLimiterByLocal(combineKey, time, count);
//                return;
//            }

            Long number = (Long) redisTemplate.execute(limitScript, keys, count, time);
            if(StringUtils.isNull(number) || number.intValue() > count){
                throw new ServiceException("访问过于频繁，请稍后重试");
            }
            log.info("限流请求'{}', 当前请求'{}', 缓存key'{}'", count, number.intValue(), key);

        } catch (ServiceException e) {
            throw e;
        } catch (Exception e) {
            throw new RuntimeException("服务限流异常，请稍后重试");
        }
    }


    /**
     * redis宕机采用本地限流
     * Google开源工具包Guava提供了限流工具类RateLimiter
     *
     * @param combineKey
     * @param time
     * @param count
     */
    private void rateLimiterByLocal(String combineKey, int time, int count) {
        // Guava提供了限流工具类RateLimiter
        com.google.common.util.concurrent.RateLimiter rateLimiter = null;
        double permitsPerSecond = count / time;

        if (!limiterMap.containsKey(combineKey)) {
            //创建令牌桶
            rateLimiter = com.google.common.util.concurrent.RateLimiter.create(permitsPerSecond);
            limiterMap.put(combineKey, rateLimiter);
            log.info("新建了令牌桶-'{}', 容量-'{}'", combineKey, permitsPerSecond);
        }
        rateLimiter = limiterMap.get(combineKey);

        boolean acquire = rateLimiter.tryAcquire();
        if (!acquire) {
            log.error("令牌桶-'{}'，获取令牌失败", combineKey);
            throw new ServiceException("访问过于频繁，请稍后重试");
        }
        log.info("令牌桶-'{}, 容量-'{}'", combineKey, permitsPerSecond);
    }


    private String getCombineKey(RateLimiter rateLimiter, JoinPoint joinPoint) {
        StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());
        if (rateLimiter.limitType() == LimitType.IP) {
            stringBuffer.append(IPUtils.getIpAddr(ServletUtils.getRequest()));
        }
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        stringBuffer.append("_").append(targetClass.getName()).append("_").append(method.getName());
        return stringBuffer.toString();
    }

    /**
     * 获取注解
     *
     * @param joinPoint
     * @return
     */
    private RateLimiter getAnnotationRateLimiter(JoinPoint joinPoint) {
        Signature signature = joinPoint.getSignature();
        MethodSignature methodSignature = (MethodSignature) signature;
        Method method = methodSignature.getMethod();

        if (method != null) {
            return method.getAnnotation(RateLimiter.class);
        }
        return null;
    }
}
